import { ITensorAllocator } from "../../Core/Backends/ITensorAllocator";
import { ActionSpec } from "../Actuators/ActionSpec";
import { ActionBuffers } from "../Actuators/IActionReceiver";
import { Utilities } from "../Utilities";
import { IApplier } from "./TensorApplier";
import { TensorProxy } from "./TensorProxy";
import { Multinomial } from "./Utils/Multinomial";



export class ContinuousActionOutputApplier implements IApplier {
    private readonly _actionSpec: ActionSpec;

    constructor(actionSpec: ActionSpec) {
        this._actionSpec = actionSpec;
    }

    public apply(tensorProxy: TensorProxy, actionIds: number[], lastActions: { [agentId: number]: ActionBuffers }): void {
        const actionSize = tensorProxy.shape[tensorProxy.shape.length - 1];
        let agentIndex = 0;
        for (let i = 0; i < actionIds.length; i++) {
            const agentId = actionIds[i];
            if (Reflect.has(lastActions, agentId)) {
                let actionBuffer = lastActions[agentId];
                if (actionBuffer.isEmpty()) {
                    actionBuffer = new ActionBuffers(this._actionSpec);
                    lastActions[agentId] = actionBuffer;
                }
                const continuousBuffer = actionBuffer.continuousActions;
                for (let j = 0; j < actionSize; j++) {
                    continuousBuffer.setValue(j,tensorProxy.data.getCache([agentIndex, j]));
                }
            }
            agentIndex++;
        }
    }
}

export class DiscreteActionOutputApplier implements IApplier {
    private readonly _actionSpec: ActionSpec;

    constructor(actionSpec: ActionSpec, seed: number, allocator: ITensorAllocator) {
        this._actionSpec = actionSpec;
    }

    public apply(tensorProxy: TensorProxy, actionIds: number[], lastActions: { [agentId: number]: ActionBuffers }): void {
        let agentIndex = 0;
        const actionSize = tensorProxy.shape[tensorProxy.shape.length - 1];
        for (let i = 0; i < actionIds.length; i++) {
            const agentId = actionIds[i];
            if (Reflect.has(lastActions, agentId)) {
                let actionBuffer = lastActions[agentId];
                if (actionBuffer.isEmpty()) {
                    actionBuffer = new ActionBuffers(this._actionSpec);
                    lastActions[agentId] = actionBuffer;
                }
                const discreteBuffer = actionBuffer.discreteActions;
                for (let j = 0; j < actionSize; j++) {
                    discreteBuffer.setValue(j,Math.floor(tensorProxy.data.getCache([agentIndex, j])));
                }
            }
            agentIndex++;
        }
    }
}

export class LegacyDiscreteActionOutputApplier implements IApplier {
    private readonly _actionSize: number[];
    private readonly _multinomial: Multinomial;
    private readonly _actionSpec: ActionSpec;
    private readonly _startActionIndices: number[];
    private readonly _cdfBuffer: number[];

    constructor(actionSpec: ActionSpec, seed: number, allocator: ITensorAllocator) {
        this._actionSize = actionSpec.branchSizes;
        this._multinomial = new Multinomial(seed);
        this._actionSpec = actionSpec;
        this._startActionIndices = Utilities.cumSum(this._actionSize);
        const largestBranch = Math.max(...this._actionSize);
        this._cdfBuffer = new Array(largestBranch);
    }

    public apply(tensorProxy: TensorProxy, actionIds: number[], lastActions: { [agentId: number]: ActionBuffers }): void {
        let agentIndex = 0;
        for (let i = 0; i < actionIds.length; i++) {
            const agentId = actionIds[i];
            if (Reflect.has(lastActions, agentId)) {
                let actionBuffer = lastActions[agentId];
                if (actionBuffer.isEmpty()) {
                    actionBuffer = new ActionBuffers(this._actionSpec);
                    lastActions[agentId] = actionBuffer;
                }
                const discreteBuffer = actionBuffer.discreteActions;
                for (let j = 0; j < this._actionSize.length; j++) {
                    this.ComputeCdf(tensorProxy, agentIndex, this._startActionIndices[j], this._actionSize[j]);
                    discreteBuffer.setValue(j,this._multinomial.sample(this._cdfBuffer, this._actionSize[j]));
                }
            }
            agentIndex++;
        }
    }

    private ComputeCdf(logProbs: TensorProxy, batch: number, channelOffset: number, branchSize: number): void {
        let maxProb = Number.NEGATIVE_INFINITY;
        for (let cls = 0; cls < branchSize; ++cls) {
            maxProb = Math.max(logProbs.data.getCache([batch, cls + channelOffset]), maxProb);
        }
        let sumProb = 0.0;
        for (let cls = 0; cls < branchSize; ++cls) {
            sumProb += Math.exp(logProbs.data.getCache([batch, cls + channelOffset]) - maxProb);
            this._cdfBuffer[cls] = sumProb;
        }
    }
}

export class MemoryOutputApplier implements IApplier {
    private readonly _memories: { [agentId: number]: number[] };

    constructor(memories: { [agentId: number]: number[] }) {
        this._memories = memories;
    }

    public apply(tensorProxy: TensorProxy, actionIds: number[], lastActions: { [agentId: number]: ActionBuffers }): void {
        let agentIndex = 0;
        const memorySize = tensorProxy.data.width;
        for (let i = 0; i < actionIds.length; i++) {
            const agentId = actionIds[i];
            let memory = this._memories[agentId];
            if (!memory || memory.length < memorySize) {
                memory = new Array(memorySize).fill(0);
            }
            for (let j = 0; j < memorySize; j++) {
                memory[j] = tensorProxy.data.getCache([agentIndex, 0, j, 0]);
            }
            this._memories[agentId] = memory;
            agentIndex++;
        }
    }
}